import os
import argparse
from PIL import Image

def create_spritesheet(image_folder, output_path, grid_size=None, sprite_size=None, default_cols=8):
    """
    Create a sprite sheet from PNG images in a folder and its subdirectories.
    
    :param image_folder: Path to the folder containing PNG images
    :param output_path: Path where the sprite sheet will be saved
    :param grid_size: Tuple (columns, rows) defining the grid layout
    :param sprite_size: Tuple (width, height) defining the size of each sprite
    :param default_cols: Default number of columns in the sprite sheet
    """
    # Get list of PNG files in the folder and subdirectories
    image_files = []
    subdirs = []
    
    # Walk through the directory tree
    for root, dirs, files in os.walk(image_folder):
        # Filter for PNG files
        png_files = [f for f in files if f.lower().endswith('.png')]
        if png_files:
            png_files.sort()  # Sort to ensure correct order
            subdir_path = os.path.relpath(root, image_folder)
            subdirs.append(subdir_path)
            # Store full paths to files
            full_paths = [os.path.join(root, f) for f in png_files]
            image_files.append(full_paths)
    
    if not image_files:
        raise ValueError("No PNG images found in the specified folder and subdirectories")
    
    # Load all images
    all_images = []
    subdir_info = []  # To track where each subdirectory's images start and end
    
    for i, file_list in enumerate(image_files):
        images = [Image.open(f) for f in file_list]
        subdir_info.append({
            'name': subdirs[i],
            'start_index': len(all_images),
            'count': len(images)
        })
        all_images.extend(images)
    
    # Determine sprite size
    if sprite_size:
        sprite_width, sprite_height = sprite_size
        # Resize all images to the specified sprite size
        all_images = [img.resize((sprite_width, sprite_height)) for img in all_images]
    else:
        # Use the size of the first image as reference
        sprite_width, sprite_height = all_images[0].size
    
    # Determine grid size
    if grid_size:
        cols, rows = grid_size
    else:
        # Use default column count
        cols = default_cols
        # Calculate rows needed based on columns and subdirectory organization
        rows = 0
        current_row = 0
        
        for subdir in subdir_info:
            # Start a new row for each subdirectory
            if subdir != subdir_info[0]:  # Not the first subdirectory
                # Move to the start of the next row
                current_row = rows
            
            # Calculate rows needed for this subdirectory
            subdir_rows = (subdir['count'] + cols - 1) // cols  # Ceiling division
            rows = current_row + subdir_rows
    
    # Create the sprite sheet
    sheet_width = cols * sprite_width
    sheet_height = rows * sprite_height
    spritesheet = Image.new('RGBA', (sheet_width, sheet_height))
    
    # Place images in the sprite sheet
    current_row = 0
    for subdir in subdir_info:
        # Start a new row for each subdirectory
        if subdir != subdir_info[0]:  # Not the first subdirectory
            # Move to the start of the next row
            current_row = (current_row * cols + cols - 1) // cols
        
        # Place images for this subdirectory
        for i in range(subdir['count']):
            img_index = subdir['start_index'] + i
            img = all_images[img_index]
            row = current_row + (i // cols)
            col = i % cols
            spritesheet.paste(img, (col * sprite_width, row * sprite_height))
        
        # Update current_row for next subdirectory
        current_row = (current_row * cols + subdir['count'] + cols - 1) // cols
    
    # Save the sprite sheet
    spritesheet.save(output_path)
    print(f"Sprite sheet saved to {output_path}")
    print(f"Grid size: {cols} columns x {rows} rows")
    print(f"Sprite size: {sprite_width} x {sprite_height}")
    print(f"Total sprites: {len(all_images)}")
    print("Subdirectories:")
    current_row = 0
    for subdir in subdir_info:
        if subdir != subdir_info[0]:  # Not the first subdirectory
            current_row = (current_row * cols + cols - 1) // cols
        start_row = current_row
        end_row = (current_row * cols + subdir['count'] - 1) // cols
        print(f"  {subdir['name']}: rows {start_row} to {end_row}")
        current_row = end_row + 1

def main():
    parser = argparse.ArgumentParser(description="Create a sprite sheet from PNG images in a folder and subdirectories")
    parser.add_argument("input_folder", help="Path to the folder containing PNG images")
    parser.add_argument("output_file", nargs='?', help="Path where the sprite sheet will be saved (default: out_dir/spritesheet.png)")
    parser.add_argument("--cols", type=int, help="Number of columns in the sprite sheet (default: 8)")
    parser.add_argument("--rows", type=int, help="Number of rows in the sprite sheet")
    parser.add_argument("--width", type=int, help="Width of each sprite in the sheet")
    parser.add_argument("--height", type=int, help="Height of each sprite in the sheet")
    
    args = parser.parse_args()
    
    # Set default output path if not provided
    if args.output_file is None:
        # Create out_dir if it doesn't exist
        if not os.path.exists("out_dir"):
            os.makedirs("out_dir")
        args.output_file = os.path.join("out_dir", "spritesheet.png")
    
    # Set default column count
    default_cols = 8
    if args.cols:
        default_cols = args.cols
    
    grid_size = None
    sprite_size = None
    
    if args.cols and args.rows:
        grid_size = (args.cols, args.rows)
    
    if args.width and args.height:
        sprite_size = (args.width, args.height)
    
    # Ensure the output directory exists
    output_dir = os.path.dirname(args.output_file)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    create_spritesheet(args.input_folder, args.output_file, grid_size, sprite_size, default_cols)

if __name__ == "__main__":
    main()